import numpy as np
import scipy.special as sc
from chaco.api import ArrayPlotData, Plot, Legend, cbrewer, palette11, palette14
from chaco.tools.api import DragZoom, PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *
def grazingAngle(distance, altitude):
    dataGrazingAngle = {}
    dataDistance = {}
    dataGrazingAngle['不考虑地球半径'] = np.arctan2(altitude / 1e3, distance)
    dataDistance['不考虑地球半径'] = np.sqrt((altitude / 1e3) ** 2 + distance ** 2)
    earthRadius = 6378.137
    r0 = earthRadius * 4 / 3
    theta = distance / r0
    r = r0 + altitude / 1e3
    x = r * np.sin(theta)
    y = r * np.cos(theta) - r0
    dataGrazingAngle['4/3地球半径'] = np.arctan2(y, x)
    dataDistance['4/3地球半径'] = np.sqrt(x ** 2 + y ** 2)
    r0 = earthRadius
    theta = distance / r0
    r = r0 + altitude / 1e3
    x = r * np.sin(theta)
    y = r * np.cos(theta) - r0
    dataGrazingAngle['真实地球半径'] = np.arctan2(y, x)
    dataDistance['真实地球半径'] = np.sqrt(x ** 2 + y ** 2)
    return dataGrazingAngle, dataDistance


class 擦地角(HasTraits):
    altitude = Range(low=0, high=20000, value=1000)
    editor = altitude.create_editor()
    editor.mode = 'slider'
    distance = np.arange(601)
    angle = Property(Dict, observe=['altitude'])
    plot = Instance(Plot)

    def _get_angle(self):
        angle, _ = grazingAngle(self.distance, self.altitude)
        for v in angle.values():
            v *= 180 / np.pi
        return angle

    @observe("angle")
    def update_plot(self, event):
        self.plot.data.update_data(
            angleNoEarth=self.angle['不考虑地球半径'],
            angle43Earth=self.angle['4/3地球半径'],
            angleEarth=self.angle['真实地球半径'],
        )

    def _plot_default(self):
        plot = Plot(ArrayPlotData(
            distance=self.distance,
            angleNoEarth=self.angle['不考虑地球半径'],
            angle43Earth=self.angle['4/3地球半径'],
            angleEarth=self.angle['真实地球半径'],
        ))
        plot.title = 'Grazing Angle'
        plot.x_axis.title = "Distance / km"
        plot.y_axis.title = "Grazing Angle / deg"
        # plot.range2d.set_bounds((0,-5), (self.distance[-1],20))
        plot.plot(("distance", "angleNoEarth"), name='No Earth radius', color=palette11[0])
        plot.plot(("distance", "angle43Earth"), name='4/3 Earth radius', color=palette11[1])
        plot.plot(("distance", "angleEarth"), name='True earth radius', color=palette11[2])
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    traits_view = View(
        Group(
            UItem(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            )),
        Item('altitude', label='探测器海拔高度/m', style='simple', editor=editor),
    )


c1hh = -73
c2hh = 20.78
c3hh = 7.351
c4hh = 25.65
c5hh = 0.0054
c1vv = -50.79
c2vv = 25.93
c3vv = 0.7093
c4vv = 21.58
c5vv = 0.00211


def sigma0(grazingAngle, frequency, seaState):
    data = {}
    data['hh'] = c1hh + c2hh * np.log10(np.sin(grazingAngle * np.pi / 180)) + (27.5 + c3hh * grazingAngle) * np.log10(frequency) / (1 + 0.95 * grazingAngle) + c4hh * (1 + seaState) ** (1 / (2 + 0.085 * grazingAngle + 0.033 * seaState)) + c5hh * grazingAngle ** 2
    data['vv'] = c1vv + c2vv * np.log10(np.sin(grazingAngle * np.pi / 180)) + (27.5 + c3vv * grazingAngle) * np.log10(frequency) / (1 + 0.95 * grazingAngle) + c4vv * (1 + seaState) ** (1 / (2 + 0.085 * grazingAngle + 0.033 * seaState)) + c5vv * grazingAngle ** 2
    return data


class 海杂波(HasTraits):
    grazingAngleRange = np.linspace(.1, 30, 150)
    frequencyRange = np.linspace(.3, 20, 150)
    seaState = Range(low=0, high=10., value=5.)
    grazingAngle = Range(low=.1, high=30, value=3)
    frequency = Range(low=.3, high=20, value=10)
    sigma0GrazingAngle = Property(Dict, observe=['seaState', 'frequency'])
    sigma0Frequency = Property(Dict, observe=['seaState', 'grazingAngle'])

    def _get_sigma0GrazingAngle(self):
        return sigma0(self.grazingAngleRange, self.frequency, self.seaState)

    def _get_sigma0Frequency(self):
        return sigma0(self.grazingAngle, self.frequencyRange, self.seaState)

    plotGrazingAngle = Instance(Plot)
    plotFrequency = Instance(Plot)

    def _plotGrazingAngle_default(self):
        plot = Plot(ArrayPlotData(
            grazingAngleRange=self.grazingAngleRange,
            hh=self.sigma0GrazingAngle['hh'],
            vv=self.sigma0GrazingAngle['vv'],
        ))
        plot.title = 'Scattering Coefficient'
        plot.x_axis.title = "Grazing Angle / deg"
        plot.y_axis.title = "Scattering Coefficient"
        # plot.range2d.set_bounds((0, -100), (30, -20))
        plot.plot(("grazingAngleRange", "hh"), name='HH', color=palette11[0])
        plot.plot(("grazingAngleRange", "vv"), name='VV', color=palette11[1])
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    def _plotFrequency_default(self):
        plot = Plot(ArrayPlotData(
            frequencyRange=self.frequencyRange,
            hh=self.sigma0Frequency['hh'],
            vv=self.sigma0Frequency['vv'],
        ))
        plot.title = 'Scattering Coefficient'
        plot.x_axis.title = "Frequency / GHz"
        plot.y_axis.title = "Scattering Coefficient"
        # plot.range2d.set_bounds((0, -100), (20, -20))
        plot.plot(("frequencyRange", "hh"), name='HH', color=palette11[0])
        plot.plot(("frequencyRange", "vv"), name='VV', color=palette11[1])
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    @observe("sigma0GrazingAngle")
    def _update_plotGrazingAngle(self, event):
        self.plotGrazingAngle.data.update_data(
            hh=self.sigma0GrazingAngle['hh'],
            vv=self.sigma0GrazingAngle['vv'],
        )

    @observe("sigma0Frequency")
    def _update_plotFrequency(self, event):
        self.plotFrequency.data.update_data(
            hh=self.sigma0Frequency['hh'],
            vv=self.sigma0Frequency['vv'],
        )

    traits_view = View(
        Group(
            UItem(
                "plotGrazingAngle",
                editor=ComponentEditor(),
            )),
        Item('frequency', label='频率/GHz'),
        Group(UItem(
            "plotFrequency",
            editor=ComponentEditor(),
            resizable=True
        )),
        Item('grazingAngle', label='擦地角/°'),
        Item('seaState', label='海况'),
        width=1400,
        height=1000,
    )


class 探测距离(HasTraits):
    distance = np.arange(601)
    beamWidthAzimuth = Range(low=0.1, high=8, value=1)
    beamWidthElevation = Range(low=0.1, high=8, value=1)
    pulseWidth = Range(low=0.1, high=30, value=1)
    altitude = Range(low=1, high=20000, value=1000)
    editor = altitude.create_editor()
    editor.mode = 'slider'
    frequency = Range(low=.3, high=20, value=10)
    seaState = Range(low=0, high=10., value=5)
    rcsSea = Property(Dict, observe=['beamWidthAzimuth', 'beamWidthElevation', 'pulseWidth', 'altitude', 'frequency', 'seaState'])

    def _get_rcsSea(self):
        grazing, r = grazingAngle(self.distance, self.altitude)
        rcsSea = {}
        for k in grazing.keys():
            area = r[k] * 1e3 * self.beamWidthAzimuth * np.pi / 180 * np.min((r[k] * 1e3 * self.beamWidthElevation * np.pi / 180, 150 * self.pulseWidth / np.cos(grazing[k])))
            # area= r[k]*1e3 * self.beamWidthAzimuth * np.pi / 180 * r[k]*1e3 * self.beamWidthElevation * np.pi / 180
            sigma0Sea = sigma0(grazing[k], self.frequency, self.seaState)
            rcsSea[k] = {}
            rcsSea[k]['hh'] = sigma0Sea['hh'] + 10 * np.log10(area)
            rcsSea[k]['vv'] = sigma0Sea['vv'] + 10 * np.log10(area)
        return rcsSea

    plot = Instance(Plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData(
            distance=self.distance,
            rcsSeaNoEarthHH=self.rcsSea['不考虑地球半径']['hh'],
            rcsSeaNoEarthVV=self.rcsSea['不考虑地球半径']['vv'],
            rcsSea43EarthHH=self.rcsSea['4/3地球半径']['hh'],
            rcsSea43EarthVV=self.rcsSea['4/3地球半径']['vv'],
            rcsSeaEarthHH=self.rcsSea['真实地球半径']['hh'],
            rcsSeaEarthVV=self.rcsSea['真实地球半径']['vv'],
        ))
        plot.title = 'Sea Clutter'
        plot.x_axis.title = "Distance / km"
        plot.y_axis.title = "Sea Clutter / dBsm"
        # plot.range2d.set_bounds((0, -10), (600, 50))
        plot.plot(("distance", "rcsSeaNoEarthHH"), name='No Earth Radius HH', color=palette11[0])
        plot.plot(("distance", "rcsSeaNoEarthVV"), name='No Earth Radius VV', color=palette11[1])
        plot.plot(("distance", "rcsSea43EarthHH"), name='4/3 Earth Radius HH', color=palette11[2])
        plot.plot(("distance", "rcsSea43EarthVV"), name='4/3 Earth Radius VV', color=palette11[3])
        plot.plot(("distance", "rcsSeaEarthHH"), name='Earth Radius HH', color=palette11[4])
        plot.plot(("distance", "rcsSeaEarthVV"), name='Earth Radius VV', color=palette11[5])
        plot.tools.append(PanTool(plot))
        plot.tools.append(ZoomTool(plot))
        plot.tools.append(DragZoom(plot, drag_button="right"))
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    @observe("rcsSea")
    def _update_plot(self, event):
        self.plot.data.update_data(
            rcsSeaNoEarthHH=self.rcsSea['不考虑地球半径']['hh'],
            rcsSeaNoEarthVV=self.rcsSea['不考虑地球半径']['vv'],
            rcsSea43EarthHH=self.rcsSea['4/3地球半径']['hh'],
            rcsSea43EarthVV=self.rcsSea['4/3地球半径']['vv'],
            rcsSeaEarthHH=self.rcsSea['真实地球半径']['hh'],
            rcsSeaEarthVV=self.rcsSea['真实地球半径']['vv'],
        )

    straits_view = View(
        Group(UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True
        )),
        Item('beamWidthAzimuth', label='方位波束宽度/°'),
        Item('beamWidthElevation', label='俯仰波束宽度/°'),
        Item('pulseWidth', label='脉冲宽度/us'),
        Item('altitude', label='海拔高度/m', editor=editor),
        Item('frequency', label='频率/GHz'),
        Item('seaState', label='海况/级'),
        resizable=True,
        width=1200,
        height=800,
    )


class Form(HasTraits):
    擦地角 = Instance(擦地角, ())
    海杂波 = Instance(海杂波, ())
    探测距离 = Instance(探测距离, ())
    traits_view = View(
        Tabbed(
            UItem('擦地角', style='custom', label='擦地角'),
            UItem('海杂波', style='custom', label='海杂波'),
            UItem('探测距离', style='custom', label='探测距离'),
        ),
        title='水面目标雷达探测距离',
        width=1200,
        height=800,
        resizable=True,
    )


if __name__ == '__main__':
    form = Form()
    form.configure_traits()
